package com.jxw.cloudpen.web.ai;

/**
 * @author ligang
 * @create 2025/6/26 14:57
 */

import com.alibaba.fastjson.JSONObject;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionResult;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

// 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
// 初始化Ark客户端，从环境变量中读取您的API Key
public class ChatCompletionsExampleV2 {
    // 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
//    static String apiKey = System.getenv("ARK_API_KEY");

    static String apiKey = "f91b3b77-e328-425a-bddd-b261c9284755";
    // 此为默认路径，您可根据业务所在地域进行配置
    static String baseUrl = "https://ark.cn-beijing.volces.com/api/v3";
    static ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
    static Dispatcher dispatcher = new Dispatcher();
    static ArkService service = ArkService.builder().dispatcher(dispatcher).connectionPool(connectionPool).baseUrl(baseUrl).apiKey(apiKey).build();
    public static void main(String[] args) {
        System.out.println("\n----- standard request -----");
        System.out.println("\n----- streaming request -----");
        System.out.println(System.currentTimeMillis());
        Long l = System.currentTimeMillis();
        final List<ChatMessage> streamMessages = new ArrayList<>();
        final ChatMessage streamSystemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你的任务是识别用户句子的情绪表达，并根据情绪表达从给定的情绪清单中选择对应的情绪反馈，最后输出json格式的情绪反馈。\n" +
                "首先，请仔细阅读以下情绪清单：\n" +
                "<情绪清单>\n" +
                "{{快乐;\n" +
                "安心;\n" +
                "尊敬;\n" +
                "赞扬;\n" +
                "相信;\n" +
                "喜爱;\n" +
                "骄傲;\n" +
                "祝愿;\n" +
                "愤怒;\n" +
                "悲伤;\n" +
                "失望;\n" +
                "内疚;\n" +
                "思念;\n" +
                "慌张;\n" +
                "恐惧;\n" +
                "害羞;\n" +
                "哭泣;\n" +
                "烦闷;\n" +
                "憎恶;\n" +
                "贬责;\n" +
                "妒忌;\n" +
                "怀疑;\n" +
                "虚荣;\n" +
                "惊奇。}}\n" +
                "现在，请仔细阅读以下用户的句子：\n" +
                "<用户句子>\n" +
                "{{#sys.query#}}\n" +
                "请按照以下步骤进行操作：\n" +
                "1. 仔细分析用户句子的情绪表达。\n" +
                "2. 从情绪清单中选择最匹配该情绪表达的情绪反馈。\n" +
                "3. 在<情绪反馈>标签中以json格式输出你的情绪反馈，格式如下：\n" +
                "{\n" +
                "    \"情绪反馈\": \"具体的情绪名称（只输出一个情绪词）\"\n" +
                "}\n" +
                "\n" +
                "请现在开始你的分析和反馈。").build();
        final ChatMessage streamUserMessage = ChatMessage.builder().role(ChatMessageRole.USER).content("今天天气很好").build();

        streamMessages.add(streamSystemMessage);
        streamMessages.add(streamUserMessage);
        ChatCompletionRequest.ChatCompletionRequestThinking chatCompletionRequestThinking =
                new ChatCompletionRequest.ChatCompletionRequestThinking("disabled");
        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .thinking(chatCompletionRequestThinking)
                // 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
                .model("ep-20250613150613-zjcp4")
                .messages(streamMessages)
                .build();
//        service.streamChatCompletion(streamChatCompletionRequest)
//                .doOnError(Throwable::printStackTrace)
//                .blockingForEach(
//                        choice -> {
//                            if (choice.getChoices().size() > 0) {
//                                System.out.println(choice.getChoices().get(0).getMessage().getContent());
//                            }
//                        }
//                );

        ChatCompletionResult chatCompletionResult = service.createChatCompletion(streamChatCompletionRequest);
        System.out.println(JSONObject.toJSONString(chatCompletionResult));
        System.out.println(System.currentTimeMillis());
        System.out.println("耗时 is  "+(System.currentTimeMillis() -l));
//        service.shutdownExecutor();
    }
}
